package com.putprize.percy

import scala.math.{log,abs,exp,pow,min,max}

import scala.collection.mutable.{HashMap => HMap}

import org.apache.log4j.Logger

class PyotrUpdater(m:Model2, N:Int, M:Int, U:Int=1, C:Int=0,I:Int=30,
				   we:Double = 0) {
  
  private
  val model = m
  
  var count = C
  
  private
  val MAX_ITER_NUM = I
  
  private
  var N1 = 0.0
  private
  var N2 = 0.0
  
  private 
  def initCV = {
		  (0 until M).map {v => {
		    (v,new Array[Double](model.K))
		  }}.toMap
  }
  
  private
  val CV = 
    if (U == 1) Map[Int,Array[Double]]()
    else {
      initCV
    }

  //def rate = pow((1+count),-0.25)
  
  val _log = Logger.getLogger(this.getClass.getName)
  
  def extractVs(xs:Map[String,Document]) = {
    xs.values.par.flatMap(x => x.vs).toSet.toArray
  }
  
  def run_em(xs:Map[String,Document], zs:Map[String,List[Int]]) = {
    
    if (count % N == 0){
      N1 = 0.0
      N2 = 0.0
    }
    
    count += xs.size
    
    val cc = new Document(Array[Int](),Array[Float](),0)
    
    val K = model.K
    
    val countV = 
      if (U == 1){
        extractVs(xs).map(v => {(v,new Array[Double](K))}).toMap
      }
      else {
        CV
      }
    
    _log.info("Count "+count)
    _log.info("Mode "+U)
    _log.info("V "+countV.size)
    
    // 1.
    _log.info("Inference ...")
    val rs = xs.par.map(x => {
      val n = x._1
      val m = x._2
      val z = zs(n)
      val r = 
        if (we > 0) PercyWeakInfer.run_vt(model, m, z, we)
        else PyotrInfer.run_vt(model, m, z)
      (n, r)
    })
    _log.info("Inference Done")
    
    // 2.
    _log.info("Count ...")
    rs.foreach(r => {
      val n = r._1
      val phi = r._2._2
      // Count
      val doc = xs(n)
      (0 until doc.n).foreach(i => {
        val v = doc.vs(i)
        val c = doc.cs(i)
        zs(n).foreach(z => {
          countV(v)(z) += c*phi(i)(z)
        })
      })
    })
    _log.info("Count Done")
    
    val iterSum = rs.map(_._2._4).reduce(_+_)
    _log.info("Iter Num "+rs.size+" "+iterSum+" "+iterSum/rs.size)
    
    val n1 = rs.map(r => r._2._3).reduce(_+_)
    val n2 = xs.par.map(x => {
      x._2.cs.reduce(_+_)
    }).reduce(_+_)
    
    _log.info("Perplexity "+count+" "+exp(-n1/n2)+" "+n1+" "+n2)
    
    _log.info("Update Model ...")
    // Update CountV
    if (U == 1){
		(0 until M).par.foreach(v => {
			val sc1 = model.sumCountV(v)
		    val sc2 =
		      if (countV.contains(v))
		        countV(v).reduce(_+_)
		      else
		        0.0
		    val rate = 
		      if (sc1 > 0.0) max(0,min(sc2/sc1,1.0))
		      else 1.0
		    (0 until K).foreach(z => {
		        val c1 = model.getCountV(v, z)
		        val c2 = 
		          if (countV.contains(v))
		            countV(v)(z)
		          else
		            0.0
		        model.setCountV(v, z, (1-rate)*c1+c2)
		      }
		    )
		})        
    } 
    else {
      if (count % N == 0){
        (0 until M).par.foreach(v => {
          (0 until K).foreach(z => {
            model.setCountV(v, z, countV(v)(z))
            countV(v)(z) = 0.0
          })
        })
      }
    }
    
    // Update CountT
    def updateCountT = {
      	val tmpCountT = model.sumCountV
		_log.info("K "+tmpCountT.size+" "+tmpCountT.sum)
		(0 until K).foreach(z => {
		    model.setCountT(z,tmpCountT(z))
		})      
    }
    
    if (U == 1){
    	updateCountT
    } 
    else {
      if (count % N == 0){
        updateCountT
      }
    }
    
    _log.info("Update Model Done")

	N1 += n1
	N2 += n2
	_log.info("Now "+exp(-N1/N2))
	
	rs.map(nr => {
	  val n = nr._1
	  val t = nr._2._1
	  (n,t)
	}).seq.toMap
  }

}
